import socket
import threading
import struct
import pickle

class VideoServer:
    def __init__(self, host='0.0.0.0', port=9999):
        self.host = host
        self.port = port
        self.server_socket = None
        self.clients = {}  # 使用字典存储客户端信息 {socket: (addr, type)}
        self.is_running = False

    def start(self):
        """启动服务器"""
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.bind((self.host, self.port))
        self.server_socket.listen(5)
        self.is_running = True
        print(f"服务器启动在 {self.host}:{self.port}")

        accept_thread = threading.Thread(target=self.accept_connections)
        accept_thread.start()

    def accept_connections(self):
        """接受新的客户端连接"""
        while self.is_running:
            try:
                client_socket, addr = self.server_socket.accept()
                print(f"新客户端连接：{addr}")
                
                # 等待客户端发送类型标识
                try:
                    client_type = client_socket.recv(16).decode().strip()
                except:
                    client_socket.close()
                    continue

                self.clients[client_socket] = (addr, client_type)
                
                # 如果是摄像头客户端，通知所有PC客户端
                if client_type == 'camera':
                    self.broadcast_camera_info('new', addr[0], client_socket)
                
                client_thread = threading.Thread(
                    target=self.handle_client,
                    args=(client_socket,)
                )
                client_thread.start()
            except:
                break

    def handle_client(self, client_socket):
        """处理客户端连接"""
        addr, client_type = self.clients.get(client_socket, (None, None))
        if not addr:
            return

        data = b""
        payload_size = struct.calcsize("L")

        while self.is_running:
            try:
                while len(data) < payload_size:
                    data += client_socket.recv(4096)
                packed_msg_size = data[:payload_size]
                data = data[payload_size:]
                msg_size = struct.unpack("L", packed_msg_size)[0]

                while len(data) < msg_size:
                    data += client_socket.recv(4096)
                frame_data = data[:msg_size]
                data = data[msg_size:]

                # 如果是摄像头客户端，广播视频帧给所有PC客户端
                if client_type == 'camera':
                    frame = pickle.loads(frame_data)
                    self.broadcast_frame(frame, client_socket)
            except:
                break

        # 客户端断开连接
        if client_socket in self.clients:
            if client_type == 'camera':
                self.broadcast_camera_info('remove', addr[0], client_socket)
            del self.clients[client_socket]
        client_socket.close()
        print(f"客户端断开连接：{addr}")

    def broadcast_frame(self, frame, sender):
        """广播视频帧给所有PC客户端"""
        frame_data = pickle.dumps(frame)
        message_size = struct.pack("L", len(frame_data))
        
        for client_socket, (_, client_type) in self.clients.items():
            if client_type == 'pc' and client_socket != sender:
                try:
                    client_socket.sendall(message_size + frame_data)
                except:
                    continue

    def broadcast_camera_info(self, action, camera_ip, camera_socket):
        """广播摄像头信息给所有PC客户端"""
        info = pickle.dumps({
            'type': 'camera_info',
            'action': action,
            'ip': camera_ip,
            'id': str(id(camera_socket))
        })
        message_size = struct.pack("L", len(info))
        
        for client_socket, (_, client_type) in self.clients.items():
            if client_type == 'pc':
                try:
                    client_socket.sendall(message_size + info)
                except:
                    continue

    def stop(self):
        """停止服务器"""
        self.is_running = False
        for client_socket in self.clients:
            client_socket.close()
        if self.server_socket:
            self.server_socket.close()
        self.clients.clear()
        print("服务器已停止")

if __name__ == "__main__":
    server = VideoServer()
    server.start() 